import copy
import random
import string

import numpy as np

from .utils import *
from spans import Span


# 关键压缩操作是 start_mus / compress_factor
# 但不改变 delay，也就不改变 dist。
def repeat_change_spans(in_span_partitions, out_span_partitions, repeat_factor, compress_factor):
    if repeat_factor == 1 and compress_factor == 1:
        return in_span_partitions, out_span_partitions

    SortPartitionsByTraceId(in_span_partitions)
    SortPartitionsByTraceId(out_span_partitions)

    assert len(in_span_partitions) == 1
    ep_in, in_spans = list(in_span_partitions.items())[0]

    for i, in_span in enumerate(in_spans):
        x = in_span.start_mus / compress_factor
        x_tid = in_span.trace_id

        for ep_out in out_span_partitions.keys():
            # print(out_span_partitions[ep_out][i].trace_id, x_tid)
            if out_span_partitions[ep_out][i].trace_id != x_tid:
                # print(out_span_partitions[ep_out][i].trace_id, x_tid)
                print("Spans are not in order!")
                assert False
            y = x + (int(out_span_partitions[ep_out][i].start_mus) - int(in_span_partitions[ep_in][i].start_mus))
            out_span_partitions[ep_out][i].start_mus = y

        in_span_partitions[ep_in][i].start_mus = x

    SortPartitionsByTime(in_span_partitions)
    SortPartitionsByTime(out_span_partitions)

    # input()

    return in_span_partitions, out_span_partitions


def repeat_change_spans2(in_span_partitions, out_span_partitions, repeat_factor, compress_factor):
    print("Compresses but doesn't repeat or shuffle")

    # if repeat_factor <= 1 and compress_factor <= 1:
    #     return in_span_partitions, out_span_partitions

    assert len(in_span_partitions) == 1
    in_span_partitions_old = copy.deepcopy(in_span_partitions)
    out_span_partitions_old = copy.deepcopy(out_span_partitions)
    ep_in, in_spans = list(in_span_partitions_old.items())[0]

    span_inds = []
    for ind, in_span in enumerate(in_spans):
        time_order = True
        for ep_out in out_span_partitions.keys():
            out_span = out_span_partitions[ep_out][ind]
            time_order = (
                    time_order
                    and (float(in_span.start_mus) <= float(out_span.start_mus))
                    and (
                            float(out_span.start_mus) + float(out_span.duration_mus)
                            <= float(in_span.start_mus) + float(in_span.duration_mus)
                    )
            )
        if time_order:
            span_inds.append(ind)

    in_span_partitions[ep_in] = []
    for ep_out in out_span_partitions_old.keys():
        out_span_partitions[ep_out] = []

    for ind in span_inds:
        # trace_id = "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(32))
        in_span = copy.deepcopy(in_spans[ind])
        print(in_span.start_mus)
        in_span.start_mus = float(in_span.start_mus) / compress_factor
        print(in_span.start_mus)
        in_span_partitions[ep_in].append(in_span)
        # print(in_span.trace_id, in_span.sid)

        for ep_out in out_span_partitions_old.keys():
            out_span = copy.deepcopy(out_span_partitions_old[ep_out][ind])
            print(out_span.start_mus)
            out_span.start_mus = float(out_span.start_mus) / compress_factor
            print(out_span.start_mus)
            input()
            # out_span.trace_id = trace_id
            out_span_partitions[ep_out].append(out_span)
            # print(out_span.trace_id, out_span.sid)
            # input()

    return in_span_partitions, out_span_partitions


def repeat_change_spans_3(in_span_partitions, out_span_partitions, repeat_factor, compress_factor):
    print("Compresses, repeats, and shuffles")

    if repeat_factor <= 1 and compress_factor <= 1:
        return in_span_partitions, out_span_partitions

    assert len(in_span_partitions) == 1
    in_span_partitions_old = copy.deepcopy(in_span_partitions)
    out_span_partitions_old = copy.deepcopy(out_span_partitions)
    ep_in, in_spans = list(in_span_partitions_old.items())[0]

    span_inds = []
    for ind, in_span in enumerate(in_spans):
        time_order = True
        for ep_out in out_span_partitions.keys():
            out_span = out_span_partitions[ep_out][ind]
            time_order = (
                    time_order
                    and (float(in_span.start_mus) <= float(out_span.start_mus))
                    and (
                            float(out_span.start_mus) + float(out_span.duration_mus)
                            <= float(in_span.start_mus) + float(in_span.duration_mus)
                    )
            )
        if time_order:
            span_inds.append(ind)

    in_span_partitions[ep_in] = []
    for ep_out in out_span_partitions_old.keys():
        out_span_partitions[ep_out] = []

    span_inds = span_inds * repeat_factor
    random.shuffle(span_inds)
    min_start_t = min(float(in_span.start_mus) for in_span in in_spans) / compress_factor
    max_start_t = max(float(in_span.start_mus) for in_span in in_spans) / compress_factor
    start_ts = sorted([random.uniform(min_start_t, max_start_t) for _ in span_inds])
    for ind, start_t in zip(span_inds, start_ts):
        # if len(in_span_partitions[ep_in]) > 40:
        #    continue
        trace_id = "".join(
            random.choice(string.ascii_lowercase + string.digits) for _ in range(32)
        )
        in_span = copy.deepcopy(in_spans[ind])
        in_span.start_mus = float(in_span.start_mus)
        offset = start_t - in_span.start_mus
        in_span.trace_id = trace_id
        in_span.start_mus += offset
        in_span_partitions[ep_in].append(in_span)
        for ep_out in out_span_partitions_old.keys():
            out_span = copy.deepcopy(out_span_partitions_old[ep_out][ind])
            out_span.start_mus = float(out_span.start_mus)
            out_span.trace_id = trace_id
            out_span.start_mus += offset
            out_span_partitions[ep_out].append(out_span)
    return in_span_partitions, out_span_partitions


# 通过修改 cache_rate，来模拟服务的动态性（dynamism）。
"""
此函数用于模拟缓存命中的情况，通过修改缓存率 `cache_rate` 来模拟服务的动态性。
根据是否使用指数分布，随机选择部分跨度（spans）进行缓存命中处理，调整跨度的持续时间和起始时间。

:param true_assignments: 真实的跨度分配关系，字典类型。
:param in_span_partitions: 输入跨度按端点划分的分区，字典类型。
:param out_span_partitions: 输出跨度按端点划分的分区，字典类型。
:param cache_rate: 缓存命中率，浮点数。
:param exponential: 是否使用指数分布来选择跨度，布尔值，默认为 False。
:return: 处理后的真实跨度分配关系，字典类型。
"""


def create_cache_hits(true_assignments, in_span_partitions, out_span_partitions, cache_rate, exponential=False):
    # 设置随机数种子，确保结果可复现
    np.random.seed(10)

    """
    在给定的分区中查找具有指定 ID 的跨度。

    :param partition: 跨度分区，列表类型。
    :param span_id: 要查找的跨度 ID。
    :return: 找到的跨度对象，如果未找到则返回 None。
    """

    def find_span(partition, span_id):
        index = -1
        # 遍历分区中的每个跨度
        for i, span in enumerate(partition):
            if span.GetId() == span_id:
                index = i
                break

        if index != -1:
            return partition[index]

    """
    调整输入和输出跨度的持续时间和起始时间，以模拟缓存命中的影响。

    :param in_span_partitions: 输入跨度按端点划分的分区，字典类型。
    :param out_span_partitions: 输出跨度按端点划分的分区，字典类型。
    :param in_span_id: 输入跨度的 ID。
    :param cache_duration_mus: 缓存持续时间，单位为微秒。
    :param eps: 输出端点列表，列表类型。
    :param chosen_ep_number: 选中的端点编号，整数。
    """

    def adjust_spans(in_span_partitions, out_span_partitions, in_span_id, cache_duration_mus, eps, chosen_ep_number):
        trace_id = in_span_id[0]
        # 调整输入跨度的持续时间
        for ep in in_span_partitions.keys():
            for span in in_span_partitions[ep]:
                if span.GetId()[0] == trace_id:
                    span.duration_mus -= cache_duration_mus
        # 调整输出跨度的起始时间
        for i, ep in enumerate(eps):
            if i > chosen_ep_number:
                for span in out_span_partitions[ep]:
                    if span.GetId()[0] == trace_id:
                        span.start_mus -= cache_duration_mus

    """
    从给定的分区中删除具有指定 ID 的跨度。

    :param partition: 跨度分区，列表类型。
    :param span_id: 要删除的跨度 ID。
    """

    def delete_span(partition, span_id):
        index = -1
        # 遍历分区中的每个跨度
        for i, span in enumerate(partition):
            if span.GetId() == span_id:
                index = i
                break

        if index != -1:
            del partition[index]

    # 获取按顺序排列的输出端点列表
    eps = GetOutEpsInOrder(out_span_partitions)
    # 选择第一个端点
    chosen_ep_number = 0
    chosen_ep = eps[chosen_ep_number]

    # 强制使用指数分布
    exponential = True
    if exponential:
        # 指数分布的 lambda 参数
        lambda_parameter = 0.001
        # 获取输入跨度的端点
        in_ep = list(in_span_partitions.keys())[0]
        # 输入跨度的数量
        num_spans = len(in_span_partitions[in_ep])
        # 从指数分布中采样
        samples = np.random.exponential(scale=1 / lambda_parameter, size=int(cache_rate * num_spans))
        # 计算采样对应的索引
        indices = [int(sample) % num_spans for sample in samples]
        # 计算概率分布
        p = np.asarray(np.exp(-lambda_parameter * np.arange(num_spans))).astype('float64')
        p = p / np.sum(p)
        # 从索引中随机选择不重复的索引
        unique_indices = np.random.choice(np.arange(num_spans), size=int(cache_rate * num_spans), replace=False, p=p)
        # print(samples)
        # print(indices)
        # print(sorted(unique_indices))
        # print(len((unique_indices)))
        # input()

    # 获取输入跨度的端点和跨度列表
    ep_in, in_spans = list(in_span_partitions.items())[0]

    # 遍历每个输入跨度
    for i, in_span in enumerate(in_spans):
        # 生成一个 0 到 999 之间的随机整数
        random_num = random.randint(0, 999)
        # 如果当前索引在随机选择的唯一索引列表中
        if i in unique_indices:
            for ep in out_span_partitions.keys():
                if ep == chosen_ep:
                    # print("\n Before:\n")
                    # print(in_span)
                    # for ep1 in out_span_partitions.keys():
                    #     print(all_spans[true_assignments[ep1][in_span.GetId()]])
                    # 获取对应输出跨度的 ID
                    span_ID = true_assignments[ep][in_span.GetId()]
                    # 查找对应的输出跨度
                    span = find_span(out_span_partitions[ep], span_ID)
                    # 将真实分配关系标记为跳过
                    # 注意到不是删除，而是标记为 skip；skip 会被识别的
                    true_assignments[ep][in_span.GetId()] = ('Skip', 'Skip')
                    # 调整跨度的持续时间和起始时间
                    adjust_spans(in_span_partitions, out_span_partitions, in_span.GetId(), span.duration_mus, eps,
                                 chosen_ep_number)
                    # 从输出跨度分区中删除对应的跨度
                    delete_span(out_span_partitions[ep], span.GetId())
                    # print("\n After:\n")
                    # print(in_span)
                    # for ep1 in out_span_partitions.keys():
                    #     if true_assignments[ep1][in_span.GetId()] in all_spans:
                    #         x = all_spans[true_assignments[ep1][in_span.GetId()]]
                    #     else:
                    #         x = "Skip"
                    #     print(x)
                    # input()
                    break

    return true_assignments
